from moviepy.video.io.ffmpeg_tools import ffmpeg_extract_subclip
import os
import json
from concurrent.futures import ThreadPoolExecutor, as_completed


class VideoSplitter:
    """视频分割器"""

    def __init__(self, drafts_path, drafts_name, output_directory, other_source_video='', callback=None):
        """
        drafts_path: 剪映草稿目录
        drafts_name: 草稿名称
        output_directory: 输出目录
        other_source_video: 其他源视频, 指定的其他源视频, 非草稿中的源视频
        """
        self.drafts_path = drafts_path.strip()
        self.drafts_name = drafts_name.strip()
        self.output_directory = output_directory.strip()
        self.other_source_video = other_source_video.strip()
        self.callback = callback if callback is not None else print
        self.drafts_json = ""
        self.result = []
        self.get_draft_file()
        self.get_main_track_video_clip()

    def get_draft_file(self):
        if not os.path.exists(self.drafts_path):
            raise ValueError(f"草稿目录 {self.drafts_path} 不存在")

        draft_full_path = os.path.join(self.drafts_path, self.drafts_name)

        if not os.path.exists(draft_full_path):
            raise ValueError(f"没有找到草稿: {draft_full_path} 不存在")

        draft_content_file = os.path.join(draft_full_path, 'draft_content.json')

        if not os.path.exists(draft_content_file):
            raise ValueError(f"草稿目录内的 草稿文件{draft_content_file} 不存在。")
        self.drafts_json = draft_content_file

    def get_main_track_video_clip(self) -> None:
        """获取主轨道视频片段信息"""
        with open(self.drafts_json, "r", encoding="utf-8") as file:
            data = json.load(file)

        main_track = data["tracks"][0]
        videos = data['materials']['videos']

        for segment in main_track['segments']:
            material_id = segment['material_id']
            source_start = segment['source_timerange']['start'] / 1000000
            source_duration = segment['source_timerange']['duration'] / 1000000

            target_start = segment['target_timerange']['start'] / 1000000
            target_duration = segment['target_timerange']['duration'] / 1000000

            source_timerange = (source_start, source_start + source_duration)
            target_timerange = (target_start, target_start + target_duration)
            for video in videos:
                if video['id'] == material_id:
                    self.result.append({
                        'id': segment['id'],
                        'material_id': material_id,
                        'source_timerange': source_timerange,
                        'target_timerange': target_timerange,
                        'path': video['path'],
                        'material_name': video['material_name']
                    })
                    break
        # print(self.result)

    def split_video(self):
        # 检查目录是否存在,不存在则创建目录
        if not os.path.exists(self.output_directory):
            os.makedirs(self.output_directory)

        def process_video(args):
            item, counter = args
            start_time, end_time = item['source_timerange'] if not self.other_source_video else item['target_timerange']
            # print(start_time, end_time)
            video_source_path = item['path'] if not self.other_source_video else self.other_source_video
            filename_extension = item['material_name'] if not self.other_source_video else video_source_path
            # 获取原文件的扩展名
            filename = os.path.basename(filename_extension)
            _, extension = os.path.splitext(filename_extension)
            # 构建输出文件名，包含路径和序号
            output_filename = os.path.join(self.output_directory, f"{filename}-{counter}{extension}")

            ffmpeg_extract_subclip(video_source_path, start_time, end_time, targetname=output_filename)
            self.callback(f"分割完成: {filename}-{counter}{extension}")

        with ThreadPoolExecutor(max_workers=8) as executor:
            counter = 1
            tasks = [(item, counter + i) for i, item in enumerate(self.result)]
            list(executor.map(process_video, tasks))

        # 等待所有任务完成
        executor.shutdown(wait=True)
        self.callback('\n')
        self.callback(f"全部片段分割完成, 视频输出目录：{self.output_directory}, 正在打开...\n\n\n")
        os.startfile(self.output_directory)
        self.callback(f"感谢您的使用, 喜欢就给我点个赞~~   抖音号: _GaKiki_   bilibili: 悟空GK")


if __name__ == '__main__':
    d_path = r"D:\JianyingPro Drafts"
    d_dir = 'aaa'
    output = r"D:\测试"
    other_source_video = r"D:\剪映导出目录\aaa.mp4"
    # other_source_video = ''
    vs = VideoSplitter(d_path, d_dir, output, other_source_video)
    vs.split_video()
